import time
import uuid
from pprint import pprint
from typing import Dict, Any, Optional, Type, Union
from dataclasses import dataclass, asdict
from collections import defaultdict

try:
    # 尝试导入生成的protobuf代码
    from messages_pb2 import (
        MessageType, MessageHeader, DataMessage, RpcRequest, RpcResponse,
        HeartbeatMessage, ServiceRegister, ServiceEndpoint, ErrorMessage,
        BatchMessage, StreamMessage
    )
    PROTOBUF_AVAILABLE = True
except ImportError as e:
    print(f"错误: {e}")
    # 如果protobuf代码未生成，提供占位符
    PROTOBUF_AVAILABLE = False
    print("警告: Protocol Buffers 代码未生成，请运行 proto/compile.sh")


@dataclass
class SerializationStats:
    """序列化统计信息"""
    total_serialized: int = 0
    total_deserialized: int = 0
    total_bytes: int = 0
    avg_serialize_time: float = 0.0
    avg_deserialize_time: float = 0.0
    error_count: int = 0


class ObjectPool:
    """对象池 - 减少GC压力"""

    def __init__(self, factory_func, max_size: int = 100):
        self.factory = factory_func
        self.pool = []
        self.max_size = max_size
        self.created_count = 0
        self.reused_count = 0

    def get(self):
        """获取对象"""
        if self.pool:
            self.reused_count += 1
            return self.pool.pop()
        else:
            self.created_count += 1
            return self.factory()

    def put(self, obj):
        """归还对象"""
        if len(self.pool) < self.max_size:
            # 清理protobuf对象
            if hasattr(obj, 'Clear'):
                obj.Clear()
            self.pool.append(obj)

    def stats(self) -> dict:
        """获取池统计信息"""
        return {
            'pool_size': len(self.pool),
            'max_size': self.max_size,
            'created_count': self.created_count,
            'reused_count': self.reused_count,
            'hit_rate': self.reused_count / (self.created_count + self.reused_count) if (self.created_count + self.reused_count) > 0 else 0
        }


class ProtobufSerializer:
    """Protocol Buffers 序列化器"""

    def __init__(self, use_object_pool: bool = True, pool_size: int = 100):
        if not PROTOBUF_AVAILABLE:
            raise ImportError("Protocol Buffers 代码未生成，请先运行编译脚本")

        self.use_object_pool = use_object_pool
        self.stats = SerializationStats()

        # 消息类型映射
        self.message_types = {
            MessageType.DATA: DataMessage,
            MessageType.RPC_REQUEST: RpcRequest,
            MessageType.RPC_RESPONSE: RpcResponse,
            MessageType.HEARTBEAT: HeartbeatMessage,
            MessageType.SERVICE_REGISTER: ServiceRegister,
            MessageType.ERROR: ErrorMessage,
        }

        # 对象池
        self.pools = {}
        if use_object_pool:
            for msg_type, msg_class in self.message_types.items():
                self.pools[msg_type] = ObjectPool(msg_class, pool_size)

    def _get_message_object(self, msg_type: MessageType):
        """获取消息对象（从池中或新建）"""
        if self.use_object_pool and msg_type in self.pools:
            return self.pools[msg_type].get()
        else:
            return self.message_types[msg_type]()

    def _return_message_object(self, msg_type: MessageType, obj):
        """归还消息对象到池中"""
        if self.use_object_pool and msg_type in self.pools:
            self.pools[msg_type].put(obj)

    def _create_header(self, msg_type: MessageType, seq_id: int = None,
                       correlation_id: str = None, **kwargs) -> MessageHeader:
        """创建消息头"""
        header = MessageHeader()
        header.type = msg_type
        header.seq_id = seq_id or int(time.time() * 1000000)  # 微秒时间戳作为默认seq_id
        header.timestamp = int(time.time() * 1000)  # 毫秒时间戳
        header.correlation_id = correlation_id or str(uuid.uuid4())

        # 设置可选字段
        if 'source_id' in kwargs:
            header.source_id = kwargs['source_id']
        if 'target_id' in kwargs:
            header.target_id = kwargs['target_id']
        if 'metadata' in kwargs:
            header.metadata.update(kwargs['metadata'])

        return header

    def serialize_data_message(self, payload: bytes, content_type: str = "binary",
                               compressed: bool = False, **header_kwargs) -> bytes:
        """序列化数据消息"""
        start_time = time.perf_counter()

        try:
            # 获取消息对象
            msg = self._get_message_object(MessageType.DATA)

            # 设置消息头
            msg.header.CopyFrom(self._create_header(MessageType.DATA, **header_kwargs))

            # 设置消息体
            msg.payload = payload
            msg.content_type = content_type
            msg.compressed = compressed

            # 序列化
            serialized = msg.SerializeToString()

            # 更新统计
            self.stats.total_serialized += 1
            self.stats.total_bytes += len(serialized)

            # 归还对象
            self._return_message_object(MessageType.DATA, msg)

            return serialized

        except Exception as e:
            self.stats.error_count += 1
            raise
        finally:
            # 更新性能统计
            duration = time.perf_counter() - start_time
            self._update_serialize_time(duration)

    def serialize_rpc_request(self, service: str, method: str, params: bytes,
                              timeout: int = 30000, params_format: str = "json",
                              **header_kwargs) -> bytes:
        """序列化RPC请求"""
        start_time = time.perf_counter()

        try:
            msg = self._get_message_object(MessageType.RPC_REQUEST)

            msg.header.CopyFrom(self._create_header(MessageType.RPC_REQUEST, **header_kwargs))
            msg.service = service
            msg.method = method
            msg.params = params
            msg.timeout = timeout
            msg.params_format = params_format

            serialized = msg.SerializeToString()

            self.stats.total_serialized += 1
            self.stats.total_bytes += len(serialized)

            self._return_message_object(MessageType.RPC_REQUEST, msg)

            return serialized

        except Exception as e:
            self.stats.error_count += 1
            raise
        finally:
            duration = time.perf_counter() - start_time
            self._update_serialize_time(duration)

    def serialize_rpc_response(self, success: bool, result: bytes = b"",
                               error_message: str = "", error_code: int = 0,
                               result_format: str = "json", execution_time: int = 0,
                               **header_kwargs) -> bytes:
        """序列化RPC响应"""
        start_time = time.perf_counter()

        try:
            msg = self._get_message_object(MessageType.RPC_RESPONSE)

            msg.header.CopyFrom(self._create_header(MessageType.RPC_RESPONSE, **header_kwargs))
            msg.success = success
            msg.result = result
            msg.error_message = error_message
            msg.error_code = error_code
            msg.result_format = result_format
            msg.execution_time = execution_time

            serialized = msg.SerializeToString()

            self.stats.total_serialized += 1
            self.stats.total_bytes += len(serialized)

            self._return_message_object(MessageType.RPC_RESPONSE, msg)

            return serialized

        except Exception as e:
            self.stats.error_count += 1
            raise
        finally:
            duration = time.perf_counter() - start_time
            self._update_serialize_time(duration)

    def serialize_heartbeat(self, node_id: str, uptime: int, metrics: Dict[str, float] = None,
                            services: list = None, **header_kwargs) -> bytes:
        """序列化心跳消息"""
        start_time = time.perf_counter()

        try:
            msg = self._get_message_object(MessageType.HEARTBEAT)

            msg.header.CopyFrom(self._create_header(MessageType.HEARTBEAT, **header_kwargs))
            msg.node_id = node_id
            msg.uptime = uptime

            if metrics:
                msg.metrics.update(metrics)
            if services:
                msg.services.extend(services)

            serialized = msg.SerializeToString()

            self.stats.total_serialized += 1
            self.stats.total_bytes += len(serialized)

            self._return_message_object(MessageType.HEARTBEAT, msg)

            return serialized

        except Exception as e:
            self.stats.error_count += 1
            raise
        finally:
            duration = time.perf_counter() - start_time
            self._update_serialize_time(duration)

    def deserialize(self, data: bytes, expected_type: MessageType = None) -> Dict[str, Any]:
        """反序列化消息"""
        start_time = time.perf_counter()
        print(data)
        try:
            # 首先尝试解析消息头来确定类型
            if expected_type is None:
                # 尝试所有可能的消息类型
                for msg_type, msg_class in self.message_types.items():
                    try:
                        temp_msg = msg_class()
                        temp_msg.ParseFromString(data)
                        expected_type = msg_type
                        break
                    except:
                        continue

                if expected_type is None:
                    raise ValueError("无法确定消息类型")

            # 使用确定的类型进行反序列化
            msg_class = self.message_types[expected_type]
            msg = msg_class()
            msg.ParseFromString(data)

            # 转换为字典格式
            result = self._protobuf_to_dict(msg)

            self.stats.total_deserialized += 1

            return result

        except Exception as e:
            self.stats.error_count += 1
            raise
        finally:
            duration = time.perf_counter() - start_time
            self._update_deserialize_time(duration)

    def _protobuf_to_dict(self, msg) -> Dict[str, Any]:
        """将protobuf消息转换为字典"""
        from google.protobuf.json_format import MessageToDict
        return MessageToDict(msg, preserving_proto_field_name=True)

    def _update_serialize_time(self, duration: float):
        """更新序列化时间统计"""
        if self.stats.total_serialized == 1:
            self.stats.avg_serialize_time = duration
        else:
            # 移动平均
            alpha = 0.1  # 平滑因子
            self.stats.avg_serialize_time = (
                    alpha * duration + (1 - alpha) * self.stats.avg_serialize_time
            )

    def _update_deserialize_time(self, duration: float):
        """更新反序列化时间统计"""
        if self.stats.total_deserialized == 1:
            self.stats.avg_deserialize_time = duration
        else:
            alpha = 0.1
            self.stats.avg_deserialize_time = (
                    alpha * duration + (1 - alpha) * self.stats.avg_deserialize_time
            )

    def get_stats(self) -> Dict[str, Any]:
        """获取序列化统计信息"""
        stats_dict = asdict(self.stats)

        # 添加对象池统计
        if self.use_object_pool:
            pool_stats = {}
            for msg_type, pool in self.pools.items():
                pool_stats[msg_type] = pool.stats()
            stats_dict['object_pools'] = pool_stats

        # 计算性能指标
        if self.stats.avg_serialize_time > 0:
            stats_dict['serialize_rate'] = 1.0 / self.stats.avg_serialize_time
        if self.stats.avg_deserialize_time > 0:
            stats_dict['deserialize_rate'] = 1.0 / self.stats.avg_deserialize_time

        return stats_dict

    def reset_stats(self):
        """重置统计信息"""
        self.stats = SerializationStats()

    def benchmark(self, test_data: list, iterations: int = 1000) -> Dict[str, Any]:
        """性能基准测试"""
        import random

        # 重置统计
        self.reset_stats()

        # 准备测试数据
        test_messages = []
        for _ in range(iterations):
            data_item = random.choice(test_data)
            test_messages.append(data_item)

        # 序列化基准测试
        start_time = time.perf_counter()
        serialized_messages = []

        for data_item in test_messages:
            if isinstance(data_item, dict):
                # 根据数据类型选择序列化方法
                if 'service' in data_item and 'method' in data_item:
                    serialized = self.serialize_rpc_request(**data_item)
                else:
                    payload = str(data_item).encode('utf-8')
                    serialized = self.serialize_data_message(payload)
                serialized_messages.append(serialized)

        serialize_time = time.perf_counter() - start_time

        # 反序列化基准测试
        start_time = time.perf_counter()

        for serialized in serialized_messages:
            self.deserialize(serialized)

        deserialize_time = time.perf_counter() - start_time

        # 计算结果
        total_size = sum(len(msg) for msg in serialized_messages)

        return {
            'iterations': iterations,
            'serialize_time': serialize_time,
            'deserialize_time': deserialize_time,
            'serialize_rate': iterations / serialize_time,
            'deserialize_rate': iterations / deserialize_time,
            'total_size': total_size,
            'avg_message_size': total_size / iterations,
            'throughput_mbps': (total_size / (1024 * 1024)) / (serialize_time + deserialize_time),
            'stats': self.get_stats()
        }


# 使用示例
if __name__ == "__main__":
    if PROTOBUF_AVAILABLE:
        # 创建序列化器
        serializer = ProtobufSerializer()

        # 测试数据消息
        payload = b"Hello, Protocol Buffers!"
        serialized = serializer.serialize_data_message(
            payload=payload,
            content_type="text",
            source_id="client_1",
            target_id="server_1"
        )

        print(f"序列化数据: {len(serialized)} bytes")

        # 反序列化
        deserialized = serializer.deserialize(serialized)
        pprint(f"反序列化结果: {deserialized}")

        # 测试RPC请求
        rpc_serialized = serializer.serialize_rpc_request(
            service="user_service",
            method="get_user",
            params=b'{"user_id": 123}',
            timeout=5000
        )

        print(f"RPC请求: {len(rpc_serialized)} bytes")

        # 获取统计信息
        stats = serializer.get_stats()
        pprint(f"统计信息: {stats}")
    else:
        print("请先编译 Protocol Buffers 文件")